/*
 * Copyright (c) 2020 - present, Inspur Genersoft Co., Ltd.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *       http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package com.inspur.edp.jit.core.compile;

import com.inspur.edp.jit.api.compile.CompileContext;
import com.inspur.edp.jit.api.compile.CompileService;
import com.inspur.edp.jit.api.compile.JavaSourceCode;
import com.inspur.edp.jit.core.common.DirectoryManager;
import com.inspur.edp.jit.core.load.ClassLoadManager;
import com.inspur.edp.jit.core.util.JarUtil;
import com.inspur.edp.jit.core.util.JitIOException;
import org.apache.commons.io.FileUtils;
import org.apache.commons.io.IOUtils;
import org.apache.commons.lang3.StringUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.tools.Diagnostic;
import javax.tools.DiagnosticCollector;
import javax.tools.JavaCompiler;
import javax.tools.JavaFileObject;
import javax.tools.StandardJavaFileManager;
import javax.tools.ToolProvider;
import java.io.File;
import java.io.IOException;
import java.io.StringReader;
import java.util.ArrayList;
import java.util.Collection;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Set;

public class CompileServiceImpl implements CompileService {
    private static final Logger logger = LoggerFactory.getLogger(CompileServiceImpl.class);
    private final JavaCompiler compiler;

    public CompileServiceImpl() {
        logger.info("new CompileServiceImpl");
        compiler = ToolProvider.getSystemJavaCompiler();
    }

    @Override
    public void compile(CompileContext context) {
        List<String> referJars = parseClassPathList(context);

        //构造StringJavaFileObject
        List<StringJavaFileObject> javaFileObjects = new ArrayList<>();
        context.getSourceCodes().forEach(
                sourceCode ->
                        javaFileObjects.add(
                                new StringJavaFileObject(sourceCode.getFullClassName(), sourceCode.getSourceCode())
                        )
        );

        //编译
        JitJavaFileManager javaFileManager = compile(referJars, javaFileObjects);

        //获取编译结果
        List<BytesJavaFileObject> compiledClasses = getJavaFileObject(context, javaFileManager);
        //更新ClassLoader
        ClassLoadManager.put(context.getFunctionType(), context.getFunctionId(), compiledClasses);

        //更新功能包缓存
        updateFunctionCache(context, compiledClasses);
    }

    /**
     * 更新功能包缓存
     */
    private void updateFunctionCache(CompileContext context, List<BytesJavaFileObject> compiledClasses) {
        String functionDir = DirectoryManager.getJitDirectory(context.getFunctionType(), context.getFunctionId());
        DirectoryManager.createDir(functionDir);
        //更新jar包
        JarUtil.write(functionDir, context.getFunctionId(), compiledClasses);
        //更新Manifest文件
        ManifestUtil.updateManifest(context);
    }

    @Override
    public void batchCompile(List<CompileContext> contexts) {
        //按照类名分组，类名不同的分在同一组，相同的分在不同组
        List<List<CompileContext>> compileGroups = buildCompileGroups(contexts);

        logger.info("bat compile {} split to {} groups", contexts.size(), compileGroups.size());
        compileGroups.forEach(group -> {

            //合并依赖
            List<String> referJars = mergeAndParseClassPathList(group);

            //合并java文件
            List<StringJavaFileObject> javaFileObjects = new ArrayList<>(group.size());
            group.forEach(context ->
                    context.getSourceCodes().forEach(sourceCode ->
                            javaFileObjects.add(
                                    new StringJavaFileObject(sourceCode.getFullClassName(), sourceCode.getSourceCode())
                            )
                    )
            );

            //批量编译
            JitJavaFileManager javaFileManager = compile(referJars, javaFileObjects);

            group.forEach(context -> {
                List<BytesJavaFileObject> compiledClasses = getJavaFileObject(context, javaFileManager);
                //更新classLoader
                ClassLoadManager.put(context.getFunctionType(), context.getFunctionId(), compiledClasses);
                //更新功能包缓存
                updateFunctionCache(context, compiledClasses);
            });
        });
    }

    /**
     * 从JitJavaFileManager中获取编译后的字节码
     */
    private List<BytesJavaFileObject> getJavaFileObject(CompileContext context, JitJavaFileManager javaFileManager) {

        List<BytesJavaFileObject> javaFileObjects = new ArrayList<>();
        context.getSourceCodes().forEach(javaSourceCode -> {
            //内部类会被编译成多个class字节码，所以需要addAll
            javaFileObjects.addAll(javaFileManager.getCompiledJavaFileObject(javaSourceCode.getFullClassName()));
        });
        return javaFileObjects;
    }

    /**
     * 合并并解析所有的依赖
     */
    private List<String> mergeAndParseClassPathList(List<CompileContext> group) {
        Set<String> referBootSet = new HashSet<>();
        Set<String> referSet = new HashSet<>();

        group.forEach(context -> {
            if (context.getReferBoots() != null) {
                referBootSet.addAll(context.getReferBoots());
            }
            if (context.getRefers() != null) {
                referSet.addAll(context.getRefers());
            }
        });

        return parseClassPathList(new ArrayList<>(referBootSet), new ArrayList<>(referSet));
    }

    /**
     * 将编译上下文里的类按照类名分组，类名不同的分在同一组，相同的分在不同组
     */
    private List<List<CompileContext>> buildCompileGroups(List<CompileContext> contexts) {
        List<Set<String>> classNameSetGroup = new ArrayList<>();
        List<List<CompileContext>> compileListGroup = new ArrayList<>();
        for (CompileContext context : contexts) {
            boolean needNewGroup = true;
            for (int i = 0; i < classNameSetGroup.size(); i++) {
                Set<String> classNameSet = classNameSetGroup.get(i);
                boolean isDistinct = true;
                for (JavaSourceCode javaSourceCode : context.getSourceCodes()) {
                    if (classNameSet.contains(javaSourceCode.getFullClassName())) {
                        isDistinct = false;
                        break;
                    }
                }
                if (isDistinct) {
                    for (JavaSourceCode javaSourceCode : context.getSourceCodes()) {
                        classNameSet.add(javaSourceCode.getFullClassName());
                    }

                    compileListGroup.get(i).add(context);
                    needNewGroup = false;
                }
            }
            if (needNewGroup) {
                Set<String> classNameSet = new HashSet<>();
                List<CompileContext> compileContexts = new ArrayList<>();

                for (JavaSourceCode javaSourceCode : context.getSourceCodes()) {
                    classNameSet.add(javaSourceCode.getFullClassName());
                }
                compileContexts.add(context);

                classNameSetGroup.add(classNameSet);
                compileListGroup.add(compileContexts);

            }
        }
        return compileListGroup;
    }

    /**
     * 解析所有的依赖jar包，将其转换为jar包的绝对路径列表。
     *
     * @param context 编译上下文
     * @return context的依赖jar包的绝对路径
     */
    private List<String> parseClassPathList(CompileContext context) {
        List<String> refers = context.getRefers();
        List<String> referBoots = context.getReferBoots();
        return parseClassPathList(referBoots, refers);
    }

    /**
     * 解析所有的依赖jar包，将其转换为jar包的绝对路径列表。
     */
    private List<String> parseClassPathList(List<String> referBoots, List<String> refers) {
        List<String> referJars = new ArrayList<>();
        //jar包解压
        if (referBoots != null && referBoots.size() > 0) {
            //从3rd目录下寻找第三方jar包
            referJars.addAll(findIn3rdDir(referBoots));
        }
        if (refers != null && refers.size() > 0) {
            refers.forEach(name -> {
                referJars.add(DirectoryManager.getHomeDir() + File.separator + name);
            });
        }
        return referJars;
    }

    /**
     * 从3rd中找到依赖的jar并返回绝对路径
     */
    private Collection<String> findIn3rdDir(List<String> referBoots) {
        List<String> refers = new ArrayList<>(referBoots.size());
        String thirdDir = DirectoryManager.get3rdDirectory();
        Collection<File> jarFiles = FileUtils.listFiles(new File(thirdDir), new String[]{"jar"}, false);
        List<String> refersBootsCopy = new ArrayList<>(referBoots);
        for (File file : jarFiles) {
            if (refersBootsCopy.isEmpty()) {
                break;
            }
            Iterator<String> iterator = refersBootsCopy.iterator();
            while (iterator.hasNext()) {
                String referBoot = iterator.next();
                if (file.getName().startsWith(referBoot)) {
                    refers.add(thirdDir + File.separator + file.getName());
                    iterator.remove();
                    break;
                }
            }
        }
        if (!refersBootsCopy.isEmpty()) {
            StringBuilder builder = new StringBuilder("未在目录中找到依赖的第三方jar包：");
            builder.append(thirdDir);
            for (String refer : referBoots) {
                builder.append(",").append(refer);
            }
            throw new JitCompileException(builder.toString());
        }
        return refers;
    }

    private JitJavaFileManager compile(List<String> refers, List<? extends JavaFileObject> javaFiles) {
        DiagnosticCollector<JavaFileObject> diagnosticCollector = new DiagnosticCollector<>();
        // 获取标准的Java文件管理器实例
        StandardJavaFileManager standardJavaFileManager = compiler.getStandardFileManager(diagnosticCollector, null, null);
        // 初始化自定义Java文件管理器实例
        JitJavaFileManager jitJavaFileManager = new JitJavaFileManager(standardJavaFileManager);
        String classPath = StringUtils.join(refers, File.pathSeparator);

        List<String> options = new ArrayList<>();
        options.add("-classpath");
        options.add(classPath);

        logger.info("compile classPath: {}", classPath);
        boolean isSuccess = compiler.getTask(null, jitJavaFileManager, diagnosticCollector, options, null, javaFiles).call();
        logger.info("compile end: {}", isSuccess);

        try {
            jitJavaFileManager.close();
        } catch (IOException e) {
            throw new JitIOException(e.getMessage(), e);
        }
        if (!isSuccess) {
            List<Diagnostic<? extends JavaFileObject>> compileError = diagnosticCollector.getDiagnostics();
            String errorInfo = getCompileErrorInfo(compileError);
            logCompileErrorClassSourceCode(compileError);
            throw new JitCompileException("编译Java文件出错:" + errorInfo);
        }
        return jitJavaFileManager;
    }

    private String getCompileErrorInfo(List<Diagnostic<? extends JavaFileObject>> compileError) {
        Map<String, List<String>> javaLinesMap = new HashMap<>();
        StringBuilder compileErrorRes = new StringBuilder();
        for (Diagnostic<? extends JavaFileObject> diagnostic : compileError) {
            compileErrorRes.append("lineNumber:").append(diagnostic.getLineNumber()).append(System.lineSeparator());
            compileErrorRes.append("columnNumber:").append(diagnostic.getColumnNumber()).append(System.lineSeparator());
            try {
                Object source = diagnostic.getSource();
                if (source instanceof StringJavaFileObject) {
                    String className = ((StringJavaFileObject) source).getClassName();
                    String sourceCode = ((StringJavaFileObject) source).getSourceCode();
                    compileErrorRes.append("className:").append(className).append(System.lineSeparator());
                    if (!javaLinesMap.containsKey(className)) {
                        List<String> lines = IOUtils.readLines(new StringReader(sourceCode));
                        javaLinesMap.put(className, lines);
                    }
                    List<String> lines = javaLinesMap.get(className);
                    long lineNumber = diagnostic.getLineNumber() - 1;//错误行号从1开始，减去1才是真实行号
                    if (lineNumber >= 0 && lineNumber < lines.size()) {
                        compileErrorRes.append("line:").append(lines.get((int) lineNumber)).append(System.lineSeparator());
                    }
                }
            } catch (Exception e) {
                logger.error("提取【" + diagnostic.getSource() + "】中的编译错误信息出错：" + e.getMessage(), e);
            }
            compileErrorRes.append("message:").append(diagnostic.getMessage(null)).append(System.lineSeparator());
        }
        return compileErrorRes.toString();
    }

    private void logCompileErrorClassSourceCode(List<Diagnostic<? extends JavaFileObject>> compileError) {
        Set<String> classNameSet = new HashSet<>();
        StringBuilder compileErrorRes = new StringBuilder();
        for (Diagnostic<? extends JavaFileObject> diagnostic : compileError) {
            Object source = diagnostic.getSource();
            if (source instanceof StringJavaFileObject) {
                String className = ((StringJavaFileObject) source).getClassName();
                String sourceCode = ((StringJavaFileObject) source).getSourceCode();
                if (!classNameSet.contains(className)) {
                    classNameSet.add(className);
                    compileErrorRes.append("类名：").append(className).append(System.lineSeparator());
                    compileErrorRes.append(sourceCode).append(System.lineSeparator());
                }
            }
        }
        logger.error("编译Java文件出错，涉及的类源码有：" + System.lineSeparator() + compileErrorRes);
    }

}
